// Copyright 2022 jeff.li. and/or its affiliates.

#include <tbir/pipeline/jit_op.h>

#include <tbir/pipeline/tx_session.h>
#include <tbir/runtime/c_runtime_api.h>
#include <tbir/runtime/container_private.h>
#include <tbir/runtime/logging.h>

namespace tbir::runtime {

    TBIR_REGISTER_NATIVE_OP(JitOp).SetThreadSafety(false);

    void JitOp::Init() {
        main_func_name_ = GetAttr<String>("main_func_name");
        jit_object_name_ = GetAttr<String>("jit_object_name");
        jit_object_ = belong_to_->FindJitObject(jit_object_name_);
        MXCHECK(jit_object_ != nullptr);
        name_ = jit_object_->PyObjectName() + "_" + name_;
        sub_ops_ = {jit_object_};
        auto pf = jit_object_->GetFunction(main_func_name_);
        forward_func_ = pf.first;
        func_meta_ = pf.second;
        self_ = jit_object_->self();
        MXCHECK(forward_func_ != nullptr && func_meta_ != nullptr)
            << "[JitOp] function not found, name: " << main_func_name_;
        MXCHECK(self_.defined()) << "[JitOp] self is nullptr";
    }

    int JitOp::Bundle(string_view folder) {
        return 0;
    }

    RTValue JitOp::Process(PyArgs inputs) const {
        PyArgs real_inputs = inputs;
        std::unique_ptr<RTView[]> new_inputs;
        if (jit_object_->options_.is_class) {
            if (func_meta_->bound_self) {
                new_inputs.reset(new RTView[inputs.size() + 1]);
                new_inputs[0] = jit_object_->self_;
                for (size_t i = 0; i < inputs.size(); ++i) {
                    new_inputs[i + 1] = inputs[i].As<RTView>();
                }
                real_inputs = PyArgs(new_inputs.get(), inputs.size() + 1);
            }
        } else {
            new_inputs.reset(new RTView[inputs.size() + 1]);
            // bundle session handler
            for (size_t i = 0; i < inputs.size(); ++i) {
                new_inputs[i] = inputs[i].As<RTView>();
            }
            new_inputs[inputs.size()] = RTView(belong_to_);
            real_inputs = PyArgs(new_inputs.get(), inputs.size() + 1);
        }
        return forward_func_(real_inputs);
    }

    RTValue JitOp::generic_call_attr(string_view func_name, PyArgs args) {
        return jit_object_->generic_call_attr(func_name, args);
    }

    String JitOp::GetHumanName(bool with_debug_info) const {
        String human_name;
        auto jit_obj_name = jit_object_->PyObjectName();
        if (jit_object_->options_.is_class) {
            String method_name = FunctionNameRules::remove_class_prefix(jit_obj_name, main_func_name_);
            human_name = jit_obj_name + "." + method_name;
        } else {
            human_name = jit_obj_name;
        }
        if (with_debug_info && jit_object_->options_.py_source_line_ >= 0) {
            auto filename = FileUtil::GetFileBasename(jit_object_->options_.py_source_file_);
            auto fileline = std::to_string(jit_object_->options_.py_source_line_);
            human_name.append(" @");
            human_name.append(filename.data(), filename.size());
            human_name.append(":");
            human_name.append(fileline.data(), fileline.size());
        }
        return human_name;
    }

}  // namespace tbir::runtime
